# -*- coding: utf-8 -*-
'''
Created on Aug 6, 2019

@author: yl
'''


from FFT_Interpolation import *
from mpl_toolkits.mplot3d import Axes3D
from scipy import signal
from scipy.optimize import curve_fit
from scipy.signal import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,
                               AutoMinorLocator)
import matplotlib
from plt_style import *
colors = color_set_medium
from tqdm import tqdm
from matplotlib.ticker import ScalarFormatter
from matplotlib.gridspec import GridSpec
import matplotlib.patches as patches
fig = plt.figure()
fig.set_size_inches(10, 3, True)
gs = GridSpec(1, 2, width_ratios=[3, 1])
ax1 = fig.add_subplot(gs[0])
ax2 = fig.add_subplot(gs[1])

j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]

pix_size = 5.3e-6
pix_num = 200
screen_diameter = pix_num * pix_size


dx = np.linspace(0, (pix_num-1)*pix_size, num=pix_num)
dy = dx
V_r_x, V_r_y, V_r_z = 0.00, 0.00, 1 # Reference vector
L, D = 0.1, 0.1  # D = displacement, L is static
V_m_x, V_m_y, V_z_m = 0.000, 0.000, 1 # Measurement vector

X, Y = np.meshgrid(dx, dy)

def f(X, Y):
    '''
        Window 2d
    '''
    window = signal.gaussian(pix_num, std=pix_num/8)
    window_x, window_y = np.meshgrid(window, window)
    window_2d = window_x*window_y
    
    k = 2 * np.pi / Lamda
    I_0 = 512
    
    '''
        Wave Generating
    '''
    diff_Z_p = D * 2 / (1-V_m_x**2-V_m_y**2) + L * 2 * ((V_m_x**2+V_m_y**2)-(V_r_x**2+V_r_y**2)) / (1-V_m_x**2-V_m_y**2) / (1-V_r_x**2-V_r_y**2)
#     d_ref = X*2*V_r_x*(1-V_r_x**2-V_r_y**2)/(1-V_r_x**2)/(1+V_r_x**2+V_r_y**2) + Y*2*V_r_y*(1-V_r_x**2-V_r_y**2)/(1-V_r_y**2)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
#     d_mea = X*2*V_m_x*(1-V_m_x**2-V_m_y**2)/(1-V_m_x**2)/(1+V_m_x**2+V_m_y**2) + Y*2*V_m_y*(1-V_m_x**2-V_m_y**2)/(1-V_m_y**2)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_mea = (X*2*V_m_x+Y*2*V_m_y)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_ref = (X*2*V_r_x+Y*2*V_r_y)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
    diff_phi = k * (diff_Z_p + (d_mea-d_ref))
    I_beat = I_0*(1+np.cos(diff_phi))
#     I_beat = I_beat.astype(np.int)
    
    '''
        Wave Generating (2nd way)
    '''
    D_d = D + X * (V_m_x - V_r_x) + Y * (V_m_y - V_r_y)
    L_d = L + X * V_r_x + Y * V_r_y
    diff_L_d_2 = D_d + (L_d + D_d) * (1 - (V_m_x**2 + V_m_y**2)) / (1 + (V_m_x**2 + V_m_y**2)) - L_d * (1 - (V_r_x**2 + V_r_y**2)) / (1 + (V_r_x**2 + V_r_y**2))
    Z = I_0 * (1 + np.cos(2 * np.pi * diff_L_d_2 / Lamda))
    
    return (I_beat) * window_2d

laser = ax1.add_patch(patches.Rectangle((5, 50),15,10,facecolor=colors['blue'],fill=True))
BS = ax1.plot([30,45],[41,66],linewidth=1,color=colors['blue'])
Ref_M = ax1.plot([30,45],[90,90],linewidth=3,color=colors['blue'],zorder=100)
Mea_M, = ax1.plot([60,60],[40,65],linewidth=3,color=colors['blue'])
Camera = ax1.add_patch(patches.Rectangle((30, 10),15,10,facecolor=colors['blue'],fill=True))
ax1.plot([38,38],[20,90],color=colors['red'])
light, = ax1.plot([20,60],[55,55],color=colors['red'])

ax1.text(0.08, 0.65, 'Laser', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
ax1.text(0.3, 0.95, 'Ref. Mirror', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
ax1.text(0.3, 0.05, 'Detector', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
ax1.text(0.7, 0.3, 'Mea. Mirror', fontsize=12, horizontalalignment='left', verticalalignment='center', transform=ax1.transAxes)
 
# note = ax1.text(0.55, 0.75, 'Displacement', fontsize=15, horizontalalignment='left', verticalalignment='center', color=color_set_normal['red'], transform=ax1.transAxes)




# style = "Simple, tail_width=0.5, head_width=4, head_length=8"
# kw = dict(arrowstyle=style, color="k")
# arrow = ax1.add_patch(patches.FancyArrowPatch((20, 20), (40, 40), connectionstyle="arc3,rad=.5", **kw))
        

ax1.set_xlim(0,100)
ax1.set_ylim(0,100)
ax1.spines['bottom'].set_visible(False)
ax1.spines['left'].set_visible(False)
ax1.tick_params(axis='both',which='both',labelbottom=[],labelleft=[],length=0)

im = ax2.imshow(f(X, Y),cmap='gray', animated=True)
# ax2.colorbar(im, fraction=0.046, pad=0.04)
line, = ax2.plot((-f(X, Y)[pix_num//2]/(5000//pix_num)+pix_num-10), color='white')
line_ver, = ax2.plot(f(X, Y)[:,pix_num//2]/(5000//pix_num)+10, np.arange(pix_num), color='white')
num = 0 
unit = 50

def updatefig(*args):
    global D, num, V_m_x, V_m_y, V_r_x, V_r_y, L
    if num<unit:
        D += -Lamda / 16
        D_layout = (0.1-D)/Lamda*10+60

        Mea_M.set_xdata([D_layout,D_layout])
        light.set_xdata([20,D_layout])
        text = 'Displacement'
        txt_color = colors['red']
        
    if num >= unit*1 and num < unit*2.5:
        V_m_y += 0.5e-4
        D_layout = (0.1-D)/Lamda*10+60
        D_4_V = V_m_y*1e3    
        Mea_M.set_xdata([D_layout-D_4_V,D_layout+D_4_V])
        text = 'Vertical Rotation'
        txt_color = colors['red']
    if num >= unit*2.5 and num < unit*4:
        V_m_x += 0.5e-4
        text = 'Horizontal Rotation'
        txt_color = colors['red']
    if num >=unit*4 and num <unit*5:
        D += Lamda / 16
        D_layout = (0.1-D)/Lamda*10+60
        D_4_V = V_m_y*1e3  
        Mea_M.set_xdata([D_layout-D_4_V,D_layout+D_4_V])
        light.set_xdata([20,D_layout])
        text = 'Displacement (tilted)'
        txt_color = colors['red']
#     V_m_x = 0 - 0.00005 * num
#     V_m_y = 0 - 0.00005 * num
    im.set_array(f(X, Y))
    line.set_ydata(-f(X, Y)[pix_num//2]/(5000//pix_num)+pix_num-10)
    line_ver.set_xdata(f(X, Y)[:,pix_num//2]/(5000//pix_num)+10)
#     note.set_text(text)
#     note.set_color(txt_color)
    
    num = num+1
    return im, line, line_ver,light, Mea_M

ani = animation.FuncAnimation(fig, updatefig, interval=50, save_count=250, blit=True)
plt.show()
# writer = animation.PillowWriter(fps=10)
ani.save(r'C:\Users\yu03\Desktop\test1.gif', writer=animation.PillowWriter(fps=20))
# ani.save(r'C:\Users\yu03\Desktop\test.mp4', writer=animation.writers['ffmpeg'])
# ani.save('test.gif')